package com.jsean.core.repository;

import java.sql.*;
import java.util.HashMap;
import java.util.Map;

public class H2OffsetRepository implements OffsetRepository {
    private static final String DB_URL = "jdbc:h2:./store/offset";
    private static final String TABLE_NAME = "consumer_offsets";

    @Override
    public Map<String, Long> loadAll() {
        Map<String, Long> offsets = new HashMap<>();
        try (Connection conn = DriverManager.getConnection(DB_URL);
             Statement stmt = conn.createStatement()) {
            stmt.execute("CREATE TABLE IF NOT EXISTS " + TABLE_NAME + 
                       " (consumer_id VARCHAR(255) PRIMARY KEY, offset_value BIGINT)");
            ResultSet rs = stmt.executeQuery("SELECT * FROM " + TABLE_NAME);
            while (rs.next()) {
                offsets.put(rs.getString("consumer_id"), rs.getLong("offset_value"));
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return offsets;
    }

    @Override
    public void persistAll(Map<String, Long> offsets) {
        try (Connection conn = DriverManager.getConnection(DB_URL);
             PreparedStatement ps = conn.prepareStatement(
                 "MERGE INTO " + TABLE_NAME + " KEY(consumer_id) VALUES (?, ?)")) {
            for (Map.Entry<String, Long> entry : offsets.entrySet()) {
                ps.setString(1, entry.getKey());
                ps.setLong(2, entry.getValue());
                ps.addBatch();
            }
            ps.executeBatch();
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    @Override
    public long getOffset(String consumerId) {
        // 可以直接从内存中获取，因为loadAll已经加载了所有数据
        return loadAll().getOrDefault(consumerId, 0L);
    }

    @Override
    public void setOffset(String consumerId, long offset) {
        Map<String, Long> offsets = loadAll();
        offsets.put(consumerId, offset);
        persistAll(offsets);
    }
}